import pymongo
from bson.errors import InvalidDocument


class MongoPipeline:
    def __init__(self, mongo_uri, mongo_db):
        self.mongo_uri = mongo_uri
        self.mongo_db = mongo_db

    def clear_data(self, data):
        """清理多层dict中存在.或$符号的key"""

        def _clear_data(_data):
            new_data = {}
            for k, v in _data.items():
                if '.' in k or '$' in k:
                    k = k.replace('.', '-').replace('$', '-')
                if isinstance(v, dict):
                    v = _clear_data(v)
                new_data[k] = v
            return new_data

        new_data = _clear_data(data)
        return new_data

    @classmethod
    def from_crawler(cls, crawler):
        # 实例化 MongoPipeline 类
        db = cls(mongo_uri=crawler.settings.get('MONGO_URI'),
                 mongo_db=crawler.settings.get('MONGO_DB'))
        return db

    def open_spider(self, spider):
        self.client = pymongo.MongoClient(self.mongo_uri)
        self.db = self.client[self.mongo_db]

    def process_item(self, item, spider):
        name = item.__class__.__name__
        data = dict(item)
        try:
            self.db[name].insert_one(data)
        except InvalidDocument as e:
            error = str(e)
            if "must not contain '.'" in error or "must not contain '$'" in error:
                new_data = self.clear_data(data)
                self.db[name].insert_one(new_data)
            else:
                # mylogger.error(f'[pipelines error] error: {e}, data: {data}')
                print(f'[pipelines error] error: {e}, data: {data}')

        return item

    def close_spider(self, spider):
        self.client.close()
